import numpy as np
import time


class AdaBoost(object):
    def __init__(self):
        self.classfierArr = None

    def stumpClassify(self, data_set, dimen, threshVal, threshIneq):
        retArray = np.ones((data_set.shape[0],))
        if threshIneq == 'lt':
            retArray[data_set[:, dimen] <= threshVal] = -1.0
        else:
            retArray[data_set[:, dimen] > threshVal] = -1.0
        return retArray

    def buildStump(self, data_array, labels_array, D):

        m, n = data_array.shape
        numSteps = 10.0
        bestStump = {}
        bestClasEst = np.zeros((m, 1))
        minError = float('inf')
        for i in range(n):
            rangeMin = data_array[:, i].min()
            rangeMax = data_array[:, i].max()
            stepSize = (rangeMax - rangeMin) / numSteps
            for j in range(-1, int(numSteps) + 1):
                for inequal in ['lt', 'gt']:
                    threshVal = (rangeMin + float(j) * stepSize)
                    predictedVals = self.stumpClassify(data_array, i, threshVal, inequal)
                    errArr = np.ones((m,))
                    errArr[predictedVals == labels_array] = 0
                    weightedError = np.sum(D * errArr)
                    if weightedError < minError:
                        minError = weightedError
                        bestClasEst = predictedVals.copy()
                        bestStump['dim'] = i
                        bestStump['thresh'] = threshVal
                        bestStump['ineq'] = inequal
        return bestStump, minError, bestClasEst

    def train(self, data_set, labels, num_iter=40):
        data_array = np.array(data_set, np.float)
        assert len(data_array.shape) == 2, 'data_set必须是维数为2的向量'
        labels_array = np.array(labels)
        assert len(labels_array.shape) == 1
        assert data_array.shape[0] == labels_array.shape[0], '样本和标注的长度必须相等'

        labels_array[labels_array > 0] = 1
        labels_array[labels_array <= 0] = -1

        weakClassArr = []
        m, _ = data_array.shape
        D = np.ones((m,), np.float) / m
        aggClassEst = np.zeros((m,))
        for i in range(num_iter):
            bestStump, error, classEst = self.buildStump(data_array, labels_array, D)
            print('D:', D)
            alpha = float(0.5 * np.log((1.0 - error) / max(error, 1e-16)))
            bestStump['alpha'] = alpha
            weakClassArr.append(bestStump)
            print('classEst:', classEst.T)
            expon = (-1 * alpha * labels_array * classEst)
            D = D * np.exp(expon)
            D = D / D.sum()
            aggClassEst += alpha * classEst
            print('aggClassEst:', aggClassEst.T)
            aggErrors = np.sign(aggClassEst) != labels_array
            errorRate = aggErrors.sum() / m
            print('toatl error:', errorRate)
            if errorRate == 0.0:
                break
        self.classfierArr = weakClassArr
        return weakClassArr

    def predict(self, data_set):
        classfierArr = self.classfierArr
        data_array = np.array(data_set, np.float)
        assert len(data_array.shape) == 2, 'data_set必须是维数为2的向量'
        m = data_array.shape[0]
        aggClassEst = np.zeros((m, ))
        for i in range(len(classfierArr)):
            classEst = self.stumpClassify(data_array, classfierArr[i]['dim'], classfierArr[i]['thresh'], \
                                          classfierArr[i]['ineq'])
            aggClassEst += classfierArr[i]['alpha'] * classEst
            # print(aggClassEst)
        return np.sign(aggClassEst)


def loadSimpleData():
    data = np.array([[1., 2.1],
                     [2., 1.1],
                     [1.3, 1.],
                     [1., 1.],
                     [2., 1.]])
    labels = [1.0, 1.0, -1.0, -1.0, 1.0]
    return data, labels


if __name__ == '__main__':
    data_set, labels = loadSimpleData()
    ada = AdaBoost()
    ada.train(data_set, labels, 30)

    print(ada.predict([[0., 0.]]))

    print()

    print(ada.predict([[5., 5.], [0., 0.]]))

